package com.yth.utils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.ObjIntConsumer;
import java.util.stream.Collectors;

/**
 * 处理树形数据的工具类
 *
 * @author yutianhong
 * @version 1.0
 * @since 2025/2/10 15:59
 */
public class TreeUtils {

    /**
     * 深度优先遍历（前序遍历）
     *
     * @param <T>            节点类型
     * @param root           根节点
     * @param childrenGetter 获取子节点的方法
     * @param action         节点处理逻辑，接收节点和当前深度
     */
    public static <T> void depthFirstPreOrder(
            T root,
            Function<T, List<T>> childrenGetter,
            ObjIntConsumer<T> action) {
        if (root == null) {
            return;
        }

        Deque<T> nodeStack = new ArrayDeque<>();
        Deque<Integer> depthStack = new ArrayDeque<>();

        // 初始化根节点
        nodeStack.push(root);
        depthStack.push(0);

        processDepthFirstPreOrder(childrenGetter, action, nodeStack, depthStack);
    }

    /**
     * 深度优先遍历（前序遍历）
     *
     * @param <T>            节点类型
     * @param roots          根节点集合
     * @param childrenGetter 获取子节点的方法
     * @param action         节点处理逻辑，接收节点和当前深度
     */
    public static <T> void depthFirstPreOrder(
            List<T> roots,
            Function<T, List<T>> childrenGetter,
            ObjIntConsumer<T> action) {
        if (isEmpty(roots)) {
            return;
        }

        Deque<T> nodeStack = new ArrayDeque<>();
        Deque<Integer> depthStack = new ArrayDeque<>();

        // 初始化根节点
        for (int i = roots.size() - 1; i >= 0; i--) {
            nodeStack.push(roots.get(i));
            depthStack.push(0);
        }

        processDepthFirstPreOrder(childrenGetter, action, nodeStack, depthStack);
    }

    /**
     * 广度优先遍历（层次遍历）
     *
     * @param root           根节点
     * @param childrenGetter 获取子节点的方法
     * @param action         节点处理逻辑
     * @param <T>            节点类型
     */
    public static <T> void breadthFirst(T root,
                                        Function<T, List<T>> childrenGetter,
                                        Consumer<T> action) {
        if (root == null) {
            return;
        }

        Queue<T> queue = new ArrayDeque<>();
        queue.offer(root);

        processBreadthFirst(childrenGetter, action, queue);
    }

    /**
     * 广度优先遍历（层次遍历）
     *
     * @param roots          根节点集合
     * @param childrenGetter 获取子节点的方法
     * @param action         节点处理逻辑
     * @param <T>            节点类型
     */
    public static <T> void breadthFirst(List<T> roots,
                                        Function<T, List<T>> childrenGetter,
                                        Consumer<T> action) {
        if (isEmpty(roots)) {
            return;
        }

        Queue<T> queue = new ArrayDeque<>();
        for (T root : roots) {
            queue.offer(root);
        }

        processBreadthFirst(childrenGetter, action, queue);
    }


    /**
     * 从扁平列表构建树形结构
     *
     * @param nodes          扁平节点列表
     * @param idGetter       获取节点唯一标识的方法
     * @param parentIdGetter 获取父节点标识的方法
     * @param childrenSetter 设置子节点的方法
     * @param <T>            节点类型
     * @param <K>            节点标识类型
     * @return 组装好的根节点列表
     */
    public static <T, K> List<T> buildTree(List<T> nodes,
                                           Function<T, K> idGetter,
                                           Function<T, K> parentIdGetter,
                                           BiConsumer<T, List<T>> childrenSetter) {
        if (isEmpty(nodes)) {
            return Collections.emptyList();
        }

        // 创建快速索引
        Map<K, T> nodeMap = nodes.stream()
                .filter(node -> idGetter.apply(node) != null)
                .collect(Collectors.toMap(idGetter, Function.identity()));

        Map<K, List<T>> parentChildrenMap = new HashMap<>();
        List<T> roots = new ArrayList<>();

        // 构建父子关系映射
        for (T node : nodes) {
            K parentId = parentIdGetter.apply(node);

            if (isRootNode(parentId, nodeMap)) {
                roots.add(node);
            } else {
                parentChildrenMap.computeIfAbsent(parentId, k -> new ArrayList<>()).add(node);
            }
        }

        // 设置子节点
        parentChildrenMap.forEach((parentId, children) -> {
            T parent = nodeMap.get(parentId);
            if (parent != null) {
                childrenSetter.accept(parent, children);
            }
        });

        return roots;
    }

    private static <T> boolean isEmpty(Collection<T> coll) {
        return coll == null || coll.isEmpty();
    }

    private static <T> void processDepthFirstPreOrder(Function<T, List<T>> childrenGetter, ObjIntConsumer<T> action, Deque<T> nodeStack, Deque<Integer> depthStack) {
        while (!nodeStack.isEmpty()) {
            T currentNode = nodeStack.pop();
            int currentDepth = depthStack.pop();
            if (currentNode == null) {
                continue;
            }

            action.accept(currentNode, currentDepth);

            // 处理子节点
            List<T> children = childrenGetter.apply(currentNode);
            if (children == null) {
                continue;
            }
            for (int i = children.size() - 1; i >= 0; i--) {
                nodeStack.push(children.get(i));
                depthStack.push(currentDepth + 1);
            }
        }
    }

    private static <T> void processBreadthFirst(Function<T, List<T>> childrenGetter, Consumer<T> action, Queue<T> queue) {
        while (!queue.isEmpty()) {
            T node = queue.poll();
            if (node == null) {
                continue;
            }
            action.accept(node);

            List<T> children = childrenGetter.apply(node);
            if (children == null) {
                continue;
            }

            for (T child : children) {
                queue.offer(child);
            }
        }
    }

    private static <T, K> boolean isRootNode(K parentId, Map<K, T> nodeMap) {
        return parentId == null || !nodeMap.containsKey(parentId);
    }

}
